#include <Eigen/Eigen>
#include <bits/stdc++.h>
#include <geometry_msgs/PoseStamped.h>
#include <mavros_msgs/State.h>
#include <ros/ros.h>
#include <trajectory_msgs/JointTrajectoryPoint.h>

#include "utils.h"

typedef Eigen::Matrix<double, 10, 1> Eigen_Vector10d;

enum FSM_EXEC_STATE {
  NO_OFFBOARD,
  TAKE_OFF,
  TAKE_OFF_EXEC,
  HOVER,
  AVOIDANCE,
  LANDING,
  ENDING,
  WAIT_TARGET,
  GEN_NEW_TRAJ,
  REPLAN_TRAJ,
  EXEC_TRAJ,
  EMERGENCY_STOP,
  SEQUENTIAL_START
};

std::vector<std::string> FSM_EXEC_STATE_STRING = {
    "NO_OFFBOARD",     "TAKE_OFF",    "TAKE_OFF_EXEC", "HOVER",
    "AVOIDANCE",       "LANDING",     "ENDING",        "WAIT_TARGET",
    "GEN_NEW_TRAJ",    "REPLAN_TRAJ", "EXEC_TRAJ",     "EMERGENCY_STOP",
    "SEQUENTIAL_START"};

FSM_EXEC_STATE exec_state_ = NO_OFFBOARD;
mavros_msgs::State current_state;
Eigen::Vector3d current_p;
Eigen::Quaterniond current_q; // w x y z
double current_yaw = 0.0;
ros::Publisher pva_pub;
std::vector<Eigen_Vector10d> pva_traj; // x y z vx vy vz ax ay az yaw
int pva_traj_idx = 0;
double SendPvaFreq = 30;
double FSMFreq = 10;
double takeoff_h = 1.0;
double takeoff_v = 0.4;
double land_v=0.4;
double hover_land_time;
double hover_avoidance_time;
Eigen::Vector3d takeoff_goal;
double hover_start_time;
std::vector<Eigen_Vector10d> avoidance_traj;
int avoidance_traj_idx = 0;

void changeState(FSM_EXEC_STATE next_state) {
  exec_state_ = next_state;
  ROS_INFO("change state to %s", FSM_EXEC_STATE_STRING[exec_state_].c_str());
}

void stateCallback(const mavros_msgs::State::ConstPtr &msg) {
  current_state = *msg;
  if (exec_state_ != NO_OFFBOARD && current_state.mode != "OFFBOARD") {
    changeState(NO_OFFBOARD);
  }
}
void posCallback(const geometry_msgs::PoseStamped::ConstPtr &msg) {
  current_p << msg->pose.position.x, msg->pose.position.y, msg->pose.position.z;
}

void attCallback(const geometry_msgs::PoseStamped::ConstPtr &msg) {
  current_q =
      Eigen::Quaterniond(msg->pose.orientation.w, msg->pose.orientation.x,
                         msg->pose.orientation.y, msg->pose.orientation.z);
  if (current_state.mode != "OFFBOARD")
    current_yaw = CalYaw(msg->pose.orientation.x, msg->pose.orientation.y,
                         msg->pose.orientation.z, msg->pose.orientation.w);
}

void avoidanceCallBack(const trajectory_msgs::JointTrajectoryPoint::ConstPtr& msg){
  Eigen_Vector10d point;
  point.head(3) << msg->positions[0], msg->positions[1], msg->positions[2];
  point(3)=msg->velocities[0];
  point(4)=msg->velocities[1];
  point(5)=msg->velocities[2];
  point(6)=msg->accelerations[0];
  point(7)=msg->accelerations[1];
  point(8)=msg->accelerations[2];
  point(9)=msg->positions[3];
  avoidance_traj.clear();
  avoidance_traj_idx=0;
  avoidance_traj.push_back(point);
}


void takeoff_landing_traj_plan(const Eigen::Vector3d &start_p,
                               const Eigen::Vector3d &end_p, double rate,
                               double v = 1.0) {
  pva_traj.clear();
  pva_traj_idx = 0;
  std::vector<Eigen::Vector3d> takeoff_traj =
      triangle_traj_plan(start_p, end_p, rate, v);
  for (int i = 0; i < takeoff_traj.size(); i++) {
    Eigen_Vector10d point = Eigen_Vector10d::Zero();
    point.head(3) = takeoff_traj[i];
    point(9) = current_yaw;
    pva_traj.push_back(point);
  }
}

void execFSMCallback(const ros::TimerEvent &e) {
  switch (exec_state_) {
  case NO_OFFBOARD: {
    hover_start_time = ros::Time::now().toSec();

    Eigen_Vector10d point = Eigen_Vector10d::Zero();
    point.head(3) = current_p;
    point(9) = current_yaw;
    pva_traj.clear();
    pva_traj_idx = 0;
    pva_traj.push_back(point);
    if (current_state.mode == "OFFBOARD") {
      changeState(TAKE_OFF);
    }
    break;
  }
  case TAKE_OFF: {
    takeoff_goal =
        Eigen::Vector3d(current_p(0), current_p(1), current_p(2) + takeoff_h);
    takeoff_landing_traj_plan(current_p, takeoff_goal, SendPvaFreq, takeoff_v);
    changeState(TAKE_OFF_EXEC);
    break;
  }
  case TAKE_OFF_EXEC: {
    if ((current_p - takeoff_goal).norm() < 0.5) {
      changeState(HOVER);
      hover_start_time = ros::Time::now().toSec();
    }
    break;
  }
  case HOVER: {
    if (!avoidance_traj.empty() &&
        ros::Time::now().toSec() - hover_start_time > hover_avoidance_time) {
      changeState(AVOIDANCE);
    } else if (ros::Time::now().toSec() - hover_start_time > hover_land_time) {
      changeState(LANDING);
    }
    break;
  }
  case AVOIDANCE: {
    break;
  }
  case LANDING: {
    takeoff_landing_traj_plan(current_p,
                              Eigen::Vector3d(current_p(0), current_p(1), 0.0),
                              SendPvaFreq, land_v);
    changeState(ENDING);
  }
  case ENDING: {

    break;
  }
  }
}

void SendPvaTimerCallback(const ros::TimerEvent &e) {
  if (exec_state_ != AVOIDANCE) {
    if (pva_traj.empty()) {
      return;
    }
    pva_traj_idx = std::min((int)(pva_traj.size() - 1), pva_traj_idx);
    Eigen_Vector10d current_point = pva_traj[pva_traj_idx];
    setPVA(
        pva_pub,
        Eigen::Vector3d(current_point(0), current_point(1), current_point(2)),
        Eigen::Vector3d(current_point(3), current_point(4), current_point(5)),
        Eigen::Vector3d(current_point(6), current_point(7), current_point(8)),
        current_point(9));
    pva_traj_idx++;
  } else {
    if (avoidance_traj.empty()) {
      return;
    }
    avoidance_traj_idx =
        std::min((int)(avoidance_traj.size() - 1), avoidance_traj_idx);
    Eigen_Vector10d current_point = avoidance_traj[avoidance_traj_idx];
    setPVA(
        pva_pub,
        Eigen::Vector3d(current_point(0), current_point(1), current_point(2)),
        Eigen::Vector3d(current_point(3), current_point(4), current_point(5)),
        Eigen::Vector3d(current_point(6), current_point(7), current_point(8)),
        current_point(9));
    avoidance_traj_idx++;
  }
}

int main(int argc, char **argv) {
  ros::init(argc, argv, "agent_command");
  ros::NodeHandle nh("~");

  nh.getParam("FSMFreq", FSMFreq);
  nh.getParam("SendPvaFreq", SendPvaFreq);
  nh.getParam("takeoff_h", takeoff_h);
  nh.getParam("takeoff_v", takeoff_v);
  nh.getParam("land_v", land_v);
  nh.getParam("hover_avoidance_time", hover_avoidance_time);
  nh.getParam("hover_land_time", hover_land_time);

  ros::Subscriber state_sub =
      nh.subscribe<mavros_msgs::State>("/mavros/state", 1, stateCallback);
  ros::Subscriber pos_sub = nh.subscribe<geometry_msgs::PoseStamped>(
      "/global_position", 1, posCallback);
  ros::Subscriber att_sub = nh.subscribe<geometry_msgs::PoseStamped>(
      "/mavros/local_position/pose", 1, attCallback);
  ros::Subscriber avoidance_sub =
      nh.subscribe<trajectory_msgs::JointTrajectoryPoint>("/avoidance_traj", 1,
                                                          avoidanceCallBack);
  pva_pub =
      nh.advertise<trajectory_msgs::JointTrajectoryPoint>("/pva_setpoint", 1);
  ros::Timer exec_timer_ =
      nh.createTimer(ros::Duration(1.0 / FSMFreq), execFSMCallback);
  ros::Timer send_pva_timer_ =
      nh.createTimer(ros::Duration(1.0 / SendPvaFreq), SendPvaTimerCallback);
  ros::spin();
}